# coding=UTF-8
"""
Copyright Ian Ross Williams, 2012

    This file is part of ArcSecond.

    ArcSecond is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    ArcSecond is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with ArcSecond.  If not, see <http://www.gnu.org/licenses/>.
"""

import pygtk
pygtk.require('2.0')
import gtk.gdk
import numpy as np
import scipy.optimize as optimize
import scipy.stats as stats
import Image, ImageChops, ImageDraw
import thread
import time
import StringIO
from arcsecond import filters
from arcsecond import pad

class notebook_deconvolution():
    image = None            #The image to be processed
    imageProcessed = None   #The processed image
    psf = None              #The point spread function
    psfSize = 8
    useGaussianPSF = False
    gaussianStd = 3
    iterations = 10
    psfRect =None        
    pixbuf = None           #The contents of the drawing area
    width = 640             #Image size
    height = 480
    deconvolve_thread_alive=False
    alive=True
    statusbar=None
    statusbarID=None
            
    (DISPLAY_ORIGINAL,
     DISPLAY_RESULT,
     DISPLAY_PSF)=range(3)
    display = DISPLAY_ORIGINAL
    
    (METHOD_WIENER,
     METHOD_LUCYRICHARDSON,
     METHOD_MAXENT)=range(3)
    method = METHOD_WIENER
     
    def __init__(self,statusbar,id):
        self.statusbar=statusbar
        self.statusbarID=id

    def __del__(self):
        self.alive=False
        self.wait_for_thread()
        while self.deconvolve_thread_alive: time.sleep(0.1)

    def iterations_changed(self, spinbutton):
        self.iterations = int(spinbutton.get_value())
        
    def radio_display_original(self,widget):
        self.display = self.DISPLAY_ORIGINAL
        self.drawingArea.queue_draw()
        
    def radio_display_result(self,widget):
        self.display = self.DISPLAY_RESULT
        self.drawingArea.queue_draw()

    def radio_display_psf(self,widget):
        self.display = self.DISPLAY_PSF
        self.drawingArea.queue_draw()
        
    def radio_psf_size(self,widget,size):
        self.psfSize=size
        if self.psfRect:
            xm = (self.psfRect[0]+self.psfRect[2])/2
            ym = (self.psfRect[1]+self.psfRect[3])/2
            self.psfRect = (xm-size/2,ym-size/2,xm+size/2,ym+size/2)
            self.psf = self.image.crop(self.psfRect)
            self.drawingArea.queue_draw()
        else:
            self.clear_psf()
        
    def radio_wiener(self,widget):
        self.method = self.METHOD_WIENER
        
    def radio_lucyrichardson(self,widget):
        self.method = self.METHOD_LUCYRICHARDSON
        
    def radio_maxent(self,widget):
        self.method = self.METHOD_MAXENT

    def checkbuttonGaussian_callback(self, widget):
        self.useGaussianPSF = not self.useGaussianPSF
        self.drawingArea.queue_draw()

    def gaussian_changed(self, spinbutton):
        self.gaussianStd = spinbutton.get_value()
        self.drawingArea.queue_draw()
            
    def clear_psf(self):
        self.psf = None
        self.psfRect = None
        self.drawingArea.queue_draw()
        
    def apply_overlays(self, image):
        #Apply overlays to image display
        image=image.convert("RGB")
        overlay = ImageChops.duplicate(image)
        draw = ImageDraw.Draw(overlay)

        if self.psfRect:
            draw.rectangle(self.psfRect, outline=(255,255,255))

        return overlay

    # Create a new backing pixmap of the appropriate size
    def configure_event(self, widget, event):
        x, y, width, height = widget.get_allocation()
        pixmap = gtk.gdk.Pixmap(widget.window, width, height)
        pixmap.draw_rectangle(widget.get_style().white_gc,
                              True, 0, 0, width, height)
    
    # Displays the currently selected frame 
    def expose_event(self, widget, event):
        im=None
        if self.useGaussianPSF:
            psfNim=(255.*self.makeGaussian(self.psfSize,self.gaussianStd)).astype(np.int8)
            self.psf = Image.fromarray(psfNim,"L")

        if self.imageProcessed and self.display==self.DISPLAY_RESULT:
            im = self.apply_overlays(self.imageProcessed)
        elif self.image and self.display==self.DISPLAY_ORIGINAL:
            im = self.apply_overlays(self.image)
        elif self.psf and self.display==self.DISPLAY_PSF:
            xsize, ysize = self.psf.size
            xsize *= 4
            ysize *= 4
            im = self.psf.resize( (xsize,ysize) )
            
        if im:
            self.pixbuf=self.image2pixbuf(im)                
            self.drawingArea.set_size_request(im.size[0], im.size[1])
            widget.window.draw_pixbuf(None,self.pixbuf, 0, 0, 0, 0, im.size[0], im.size[1]) 

    def button_press_event(self, widget, event):
        if event.button == 1 and self.image:
            x0,y0 = int(event.x-self.psfSize/2), int(event.y-self.psfSize/2)
            x1,y1 = int(event.x+self.psfSize/2), int(event.y+self.psfSize/2)
            if x0<0 or x1>self.image.size[0]: return
            if y0<0 or y1>self.image.size[1]: return
            self.psfRect = (x0,y0,x1,y1)
            self.psf = self.image.crop(self.psfRect)
            self.drawingArea.queue_draw()

    def makeGaussian(self,size, fwhm = 3):
        """ Make a square gaussian kernel.

        size is the length of a side of the square
        fwhm is full-width-half-maximum, which
        can be thought of as an effective radius.
        """

        x = np.arange(0, size, 1, np.float32)
        y = x[:,np.newaxis]
        x0 = y0 = size // 2
        return np.exp(-4*np.log(2) * ((x-x0)**2 + (y-y0)**2) / fwhm**2)
                
    def buttonDeconvolve_callback(self, button):
        if not self.image:
            self.load_image()
            
        if self.useGaussianPSF:
            psfNim=(255.*self.makeGaussian(self.psfSize,self.gaussianStd)).astype(np.int8)
            self.psf = Image.fromarray(psfNim,"L")
            
        if not self.psf:
            md = gtk.MessageDialog(None, 
                gtk.DIALOG_DESTROY_WITH_PARENT, gtk.MESSAGE_INFO, 
                gtk.BUTTONS_CLOSE, "Please click on a star or open a Point Spread Function.")
            md.run()
            md.destroy()

        if self.image and self.psf and not self.deconvolve_thread_alive:
            thread.start_new_thread(self.thread_deconvolve, (self,))
         
    def thread_deconvolve(self,arg):
        if not self.alive: return
        self.deconvolve_thread_alive=True
        
        if self.image.mode=="RGB":
            HUE,SAT,L = filters.rgb_to_hsl(self.image)
            im = L
        else:
            im = self.image.convert("F")
            
        d = np.asarray(im, dtype=np.float)/255.
            
        psf = self.psf.convert("F")
        p = np.asarray(psf, dtype=np.float)
        p /= np.sum(p)
        
        #Pad the image data
        r1,c1=d.shape
        r2,c2=p.shape
        rows = r1+r2
        cols = c1+c2        
        if pad:
            pr2 = int(np.log(rows)/np.log(2.0) + 1.0 )
            pc2 = int(np.log(cols)/np.log(2.0) + 1.0 )
            rows = 2**pr2
            cols = 2**pc2

        dpad = np.zeros((rows,cols),dtype=np.float)
        dpad[r2/2:r2/2+r1, c2/2:c2/2+c1] = d
        d = dpad
        
        if self.method==self.METHOD_WIENER:
            d = filters.normalise(d)
            Y = np.fft.rfft2(d, s=d.shape)
            H = np.fft.rfft2(p, s=d.shape)
            S = np.fft.rfft2(d, s=d.shape)
            #N = np.random.rand(np.prod(S.shape)).reshape(S.shape)
            N = np.ones(np.prod(S.shape)).reshape(S.shape)
            underrelax=0.5
            for i in range(self.iterations):             
                NtoS = np.mean(np.abs(N)**2)/np.mean(np.abs(S)**2)
                StoN = 1./(np.max([1e-8,np.abs(NtoS)]))
                self.statusbar.pop(self.statusbarID)
                self.statusbar.push(self.statusbarID, "Signal to noise ratio = %f" % StoN )
                G = H.conj() / (np.abs(H)**2 + NtoS)
                X = G * Y
                S = np.abs(X)**2
                N = underrelax*N + (1.-underrelax)*(np.abs(Y)**2-S)
                
            u = np.fft.irfft2(X)
            
        elif self.method==self.METHOD_LUCYRICHARDSON:
            """
            http://www.ts-imaging.net/webhelp/x-tract/HID_DECONV_RICHLUCY.htm
            u_new = u . d/c correlate p
            where c = p convolve u
            """
            eps = 1.e-8
            u = np.ones(d.shape)*np.mean(d)
            for i in range(self.iterations):
                self.statusbar.pop(self.statusbarID)
                self.statusbar.push(self.statusbarID, "Iteration %d of %d" % ((i+1),self.iterations) )
                c = self.fftconvolve(u,p)
                c[c==0] = eps
                u = u*self.fftcorrelate(d/c,p)
                
        elif self.method==self.METHOD_MAXENT:
            d = filters.normalise(d)
            Y = np.fft.rfft2(d)
            H = np.fft.rfft2(p, s=d.shape)
            
            (dist, bins) = np.histogram(d*255.,np.arange(255),normed=True)
            dEntropy =  stats.entropy(dist)

            def wiener_entropy(NtoS,H):
                StoN = 1./(np.max([1e-8,np.abs(NtoS)]))
                self.statusbar.pop(self.statusbarID)
                self.statusbar.push(self.statusbarID, "Signal to noise ratio = %f" % StoN )
                G = H.conj() / (np.abs(H)**2 + np.abs(NtoS))            
                X = G * Y
                u = np.fft.irfft2(X)
                dist = 0
                (dist, bins) = np.histogram(u*255.,np.arange(255),normed=True)
                uEntropy =  stats.entropy(dist)
                deltaEntropy = uEntropy-dEntropy
                L1 = np.mean(np.abs(u-d))
                score= -deltaEntropy + 10.*L1
                return score
            
            NtoS = optimize.brent(wiener_entropy, (H,), tol=1e-4)
            G = H.conj() / (np.abs(H)**2 + np.abs(NtoS))
            X = G * Y
            u = np.fft.irfft2(X)
            
        u = u[:r1,:c1]
        umid=u[self.psfSize:-self.psfSize,self.psfSize:-self.psfSize]
        dmid=d[self.psfSize:-self.psfSize,self.psfSize:-self.psfSize]
        #u /= np.max(umid)
        u = np.clip(u,0.0,1.0)
        
        #Gamma correct the result
        #umed = np.mean(umid)
        #dmed = np.mean(dmid)
        #gamma = np.log(dmed)/np.log(umed)
        #u = u ** gamma
        
        if self.image.mode=="RGB":
            R,G,B = filters.hsl_to_rgb(HUE, SAT, 255.*u)
            R = Image.fromarray(R).convert("L")
            G = Image.fromarray(G).convert("L")
            B = Image.fromarray(B).convert("L")        
            self.imageProcessed = Image.merge("RGB", (R,G,B))
        else:
            self.imageProcessed = Image.fromarray(255.*u)
            self.imageProcessed = self.imageProcessed.convert("F")
        
        self.statusbar.pop(self.statusbarID)
        self.statusbar.push(self.statusbarID, "Done" )
        self.display=self.DISPLAY_RESULT
        self.radioButtonResult.set_active(True)            

        # Display the processed image 
        self.drawingArea.queue_draw()
        
        self.deconvolve_thread_alive=False
        
    def wait_for_thread(self):
        self.alive=False
        while self.deconvolve_thread_alive: time.sleep(0.1)

    def fftconvolve(self, f,g):
        F = np.fft.rfft2(f,s=f.shape)
        G = np.fft.rfft2(g,s=f.shape)
        return np.fft.irfft2(F*G)

    def fftcorrelate(self, f,g):
        F = np.fft.rfft2(f,s=f.shape)
        G = np.fft.rfft2(g,s=f.shape)
        return np.fft.irfft2(F*G.conj())
    
    def load_image(self):
        imageDlg = gtk.FileChooserDialog(
            title="Select an image for post-processing",
            action=gtk.FILE_CHOOSER_ACTION_OPEN,
            buttons=(gtk.STOCK_CANCEL,gtk.RESPONSE_CANCEL,gtk.STOCK_OPEN,gtk.RESPONSE_OK))
        imageDlg.set_default_response(gtk.RESPONSE_CANCEL)
        filter = gtk.FileFilter()
        filter.set_name("Images")
        filter.add_mime_type("image/png")
        filter.add_mime_type("image/ppm")
        filter.add_mime_type("image/jpeg")
        filter.add_mime_type("image/gif")
        filter.add_mime_type("image/bmp")
        filter.add_mime_type("image/psd")
        filter.add_pattern("*.png")
        filter.add_pattern("*.jpg")
        filter.add_pattern("*.gif")
        filter.add_pattern("*.bmp")
        filter.add_pattern("*.xbm")
        filter.add_pattern("*.tif")
        filter.add_pattern("*.xpm")
        filter.add_pattern("*.psd")
        imageDlg.add_filter(filter)

        filter = gtk.FileFilter()
        filter.set_name("All files")
        filter.add_pattern("*")
        imageDlg.add_filter(filter)

        response = imageDlg.run()
        if response == gtk.RESPONSE_OK:
            filename = imageDlg.get_filename()                
            imageDlg.destroy()
        else:
            imageDlg.destroy()
            return

        try:
            self.image = Image.open(filename)
        except ValueError:
            md = gtk.MessageDialog(
                None, 
                0, gtk.MESSAGE_ERROR, 
                gtk.BUTTONS_CLOSE, "The selected file was not a valid image format.")
            md.run()
            md.destroy()
            self.image = None
        except AttributeError:
            self.image = None
        else:
            self.width,self.height = self.image.size
            self.drawingArea.set_size_request(self.width, self.height)
            self.drawingArea.queue_resize()        

        self.imageProcessed = self.image
        
    def load_psf(self):
        self.clear_psf()
        imageDlg = gtk.FileChooserDialog(
            title="Select a point spread function to use for post-processing",
            action=gtk.FILE_CHOOSER_ACTION_OPEN,
            buttons=(gtk.STOCK_CANCEL,gtk.RESPONSE_CANCEL,gtk.STOCK_OPEN,gtk.RESPONSE_OK))
        imageDlg.set_default_response(gtk.RESPONSE_CANCEL)
        filter = gtk.FileFilter()
        filter.set_name("Images")
        filter.add_mime_type("image/png")
        filter.add_mime_type("image/ppm")
        filter.add_mime_type("image/jpeg")
        filter.add_mime_type("image/gif")
        filter.add_mime_type("image/bmp")
        filter.add_mime_type("image/psd")
        filter.add_pattern("*.png")
        filter.add_pattern("*.jpg")
        filter.add_pattern("*.gif")
        filter.add_pattern("*.bmp")
        filter.add_pattern("*.xbm")
        filter.add_pattern("*.tif")
        filter.add_pattern("*.xpm")
        filter.add_pattern("*.psd")
        imageDlg.add_filter(filter)

        filter = gtk.FileFilter()
        filter.set_name("All files")
        filter.add_pattern("*")
        imageDlg.add_filter(filter)

        response = imageDlg.run()
        if response == gtk.RESPONSE_OK:
            filename = imageDlg.get_filename()                
        else:
            filename = None
        
        imageDlg.destroy()

        try:
            self.psf = Image.open(filename)
        except ValueError:
            md = gtk.MessageDialog(
                None, 
                0, gtk.MESSAGE_ERROR, 
                gtk.BUTTONS_CLOSE, "The selected file was not a valid image format.")
            md.run()
            md.destroy()
            self.psf = None
        except AttributeError:
            self.psf = None

    def save_psf(self):
        im=self.psf
        if im:
            imageDlg = gtk.FileChooserDialog(
                title="Save point spread function",
                action=gtk.FILE_CHOOSER_ACTION_SAVE,
                buttons=(gtk.STOCK_CANCEL,gtk.RESPONSE_CANCEL,gtk.STOCK_SAVE,gtk.RESPONSE_OK))
            imageDlg.set_default_response(gtk.RESPONSE_CANCEL)
            filter = gtk.FileFilter()
            filter.set_name("Images")
            filter.add_mime_type("image/png")
            filter.add_mime_type("image/ppm")
            filter.add_mime_type("image/jpeg")
            filter.add_mime_type("image/gif")
            filter.add_mime_type("image/bmp")
            filter.add_mime_type("image/psd")
            filter.add_pattern("*.png")
            filter.add_pattern("*.jpg")
            filter.add_pattern("*.gif")
            filter.add_pattern("*.bmp")
            filter.add_pattern("*.xbm")
            filter.add_pattern("*.tif")
            filter.add_pattern("*.xpm")
            filter.add_pattern("*.psd")
            imageDlg.add_filter(filter)
    
            filter = gtk.FileFilter()
            filter.set_name("All files")
            filter.add_pattern("*")
            imageDlg.add_filter(filter)
    
            response = imageDlg.run()
            if response == gtk.RESPONSE_OK:
                filename = imageDlg.get_filename()
                if im.mode=="F": 
                    nim=np.asarray(im)*255
                    im=Image.fromarray(nim)
                    im=im.convert("I")
                
                try:
                    im.save(filename)
                except IOError:
                    msg= "Could not save image."
                    print msg
                    self.statusbar.pop(self.statusbarID)
                    self.statusbar.push(self.statusbarID,msg)
                
            imageDlg.destroy()

    def get_page(self):
        hbox = gtk.HBox()
        
        vbox = gtk.VBox()
        hbox.pack_start(vbox,False,False,0)
        
        #Add radio buttons to select the displayed image
        frameDisplay = gtk.Frame("Display:")
        vboxDisplay = gtk.VBox()
        radioButton = gtk.RadioButton(None, "Original")
        radioButton.set_active(self.method==self.DISPLAY_ORIGINAL)
        radioButton.connect("toggled", self.radio_display_original)
        vboxDisplay.pack_start(radioButton)
        radioButton = gtk.RadioButton(radioButton, "Result")
        radioButton.set_active(self.method==self.DISPLAY_RESULT)
        radioButton.connect("toggled", self.radio_display_result)
        self.radioButtonResult = radioButton
        vboxDisplay.pack_start(radioButton)
        radioButton = gtk.RadioButton(radioButton,"Point Spread Function")
        radioButton.set_active(self.display==self.DISPLAY_PSF)
        radioButton.connect("toggled", self.radio_display_psf)
        vboxDisplay.pack_start(radioButton)
        frameDisplay.add(vboxDisplay)
        vbox.pack_start(frameDisplay,False,False,0)
        
        #Add radio buttons to select the PSF size
        frameDisplay = gtk.Frame("PSF size:")
        vboxDisplay = gtk.VBox()
        radioButton = gtk.RadioButton(None, "8")
        radioButton.connect("toggled", self.radio_psf_size,8)
        vboxDisplay.pack_start(radioButton)
        radioButton = gtk.RadioButton(radioButton, "16")
        radioButton.connect("toggled", self.radio_psf_size,16)
        vboxDisplay.pack_start(radioButton)
        radioButton = gtk.RadioButton(radioButton, "32")
        radioButton.connect("toggled", self.radio_psf_size,32)
        vboxDisplay.pack_start(radioButton)
        radioButton = gtk.RadioButton(radioButton,"64")
        radioButton.connect("toggled", self.radio_psf_size,64)
        vboxDisplay.pack_start(radioButton)
        frameDisplay.add(vboxDisplay)
        vbox.pack_start(frameDisplay,False,False,0)  
              
        #Add radio buttons to select the method
        frameMethod = gtk.Frame("Method:")
        vboxMethod = gtk.VBox()
        radioButton = gtk.RadioButton(None,"Iterative Wiener")
        radioButton.set_active(self.method==self.METHOD_WIENER)
        radioButton.connect("toggled", self.radio_wiener)
        vboxMethod.pack_start(radioButton)
        radioButton = gtk.RadioButton(radioButton, "Lucy-Richardson")
        radioButton.set_active(self.method==self.METHOD_LUCYRICHARDSON)
        radioButton.connect("toggled", self.radio_lucyrichardson)
        vboxMethod.pack_start(radioButton)
        radioButton = gtk.RadioButton(radioButton, "Maximum Entropy")
        radioButton.set_active(self.method==self.METHOD_MAXENT)
        radioButton.connect("toggled", self.radio_maxent)
        vboxMethod.pack_start(radioButton)
        frameMethod.add(vboxMethod)
        vbox.pack_start(frameMethod,False,False,0)

        #Add a frame containing a spin widget to set the number of deconvolution iterations required
        frameIter = gtk.Frame("Number of iterations:")
        self.adjustmentIterations = gtk.Adjustment(value=self.iterations, lower=1, upper=1e6, step_incr=1, page_incr=10)
        self.spinbuttonIterations = gtk.SpinButton(self.adjustmentIterations)
        self.spinbuttonIterations.set_digits(0)
        self.spinbuttonIterations.connect("value-changed",self.iterations_changed)
        frameIter.add(self.spinbuttonIterations)
        
        #Add a frame containing a spin widget to set the size of a Gaussian kernel
        frameGaussian = gtk.Frame("Gaussian PSF:")
        vboxGaussian = gtk.VBox()
        self.adjustmentGaussian = gtk.Adjustment(value=self.gaussianStd, lower=.1, upper=1e6, step_incr=0.1, page_incr=1)
        self.spinbuttonGaussian = gtk.SpinButton(self.adjustmentGaussian)
        self.spinbuttonGaussian.set_digits(1)
        self.spinbuttonGaussian.connect("value-changed",self.gaussian_changed)
        self.checkbuttonGaussian = gtk.CheckButton("Enabled")
        self.checkbuttonGaussian.connect("clicked", self.checkbuttonGaussian_callback)
        vboxGaussian.pack_start(self.spinbuttonGaussian)
        vboxGaussian.pack_start(self.checkbuttonGaussian)
        frameGaussian.add(vboxGaussian)
        
        #Create a button to start the image registration
        self.buttonDeconvolve = gtk.Button("Deconvolve")
        self.buttonDeconvolve.set_flags(gtk.CAN_DEFAULT)
        self.buttonDeconvolve.connect("clicked", self.buttonDeconvolve_callback)
        vbox.pack_end(self.buttonDeconvolve,False,False,0)
        vbox.pack_end(frameIter, expand=False, fill=False, padding=0)
        vbox.pack_end(frameGaussian, expand=False, fill=False, padding=0)

        #Construct the capture view area 
        scrolledWindow = gtk.ScrolledWindow()
        scrolledWindow.set_policy(gtk.POLICY_AUTOMATIC, gtk.POLICY_AUTOMATIC) 
        self.drawingArea = gtk.DrawingArea()
        self.drawingArea.set_size_request(self.width, self.height)        
        scrolledWindow.add_with_viewport(self.drawingArea) 
        hbox.pack_end(scrolledWindow, expand=True, fill=True, padding=0)
        
        # Signals used to handle backing pixmap
        self.drawingArea.connect("configure_event", self.configure_event)
        self.drawingArea.connect("expose_event", self.expose_event)
    
        # Event signals
        self.drawingArea.connect("button_press_event", self.button_press_event)
    
        self.drawingArea.set_events(gtk.gdk.EXPOSURE_MASK
                                | gtk.gdk.LEAVE_NOTIFY_MASK
                                | gtk.gdk.BUTTON_PRESS_MASK
                                | gtk.gdk.POINTER_MOTION_MASK
                                | gtk.gdk.POINTER_MOTION_HINT_MASK)
        
        hbox.show_all()
        return hbox

    def image2pixbuf(self,im):    
        if im:
            file1 = StringIO.StringIO()  
            im.save(file1, "ppm")  
            contents = file1.getvalue()  
            file1.close()  
            loader = gtk.gdk.PixbufLoader("pnm")  
            loader.write(contents, len(contents))  
            pixbuf = loader.get_pixbuf()  
            loader.close()
            return pixbuf
